#include <iostream>
#include <vector>
#include "util_io.hpp"
#include "get_weight_l.hpp"
#include <omp.h>
#include <cstring>
#include <cassert>
//#define NDEBUG
using namespace std;
using namespace KMeansMultiple;

/*
 * 输入:
 *      label_l n_item * n_overlap
 *      base n_item * dimension
 * 输出:
 *      weight_l.fvecs
 * siftsmall_10_16_nwq_4_kmeans_multiple_dot_product_
 *
 * def parallel_get_weight_l(data, label_l):
    get_label_map()
    weight_l = None
    return weight_l


    def get_label_map(label_l_l, n_cluster):
        label_map = []
        n_point_label = []
        for cluster_i in range(n_cluster):
            base_idx_i = np.argwhere(label_l_l == cluster_i)[:, 0]
            label_map.append(base_idx_i)
            n_point_label.append(len(base_idx_i))
        return label_map, n_point_label
 *
 */

int main(int argc, char **argv) {
    if (argc != 6) {
        std::cout << argv[0] << " data_fname basic_dir n_classifier n_cluster n_overlap" << std::endl;
        exit(-1);
    }
    printf("data_fname %s\nbasic_dir %s\nn_classifier %s\nn_cluster %s\nn_overlap %s\n", argv[1], argv[2], argv[3],
           argv[4], argv[5]);

    char *data_fname = argv[1];
    char *basic_dir = argv[2];
    int n_classifier = atoi(argv[3]);
    int n_cluster = atoi(argv[4]);
    int n_overlap = atoi(argv[5]);

    int n_item, dimension;
    float *base = load_vecs<float>(data_fname, n_item, dimension);
    vector<vector<float>> weight_l(n_item);
    for (int i = 0; i < n_item; i++) {
        weight_l[i] = vector<float>(n_overlap);
    }
    
//#pragma omp parallel for
    for (int i = 0; i < n_classifier; i++) {
        char label_l_fname[256];
        sprintf(label_l_fname, "%s/Classifier_%d/dataset_partition/label.ivecs", basic_dir, i);
        int tmp_n_overlap;
        int *label_l = load_vecs<int>(label_l_fname, n_item, tmp_n_overlap);
        assert(tmp_n_overlap == n_overlap);
        for (int j = 0; j < n_item; j++) {
            memset(weight_l[j].data(), 0, sizeof(float) * n_overlap);
        }

        // get label_map
        vector<vector<int>> label_map = get_label_map(label_l, n_cluster, n_item, n_overlap);

//        char tmp_str[256] = "";
//        for (int j = 0; j < n_cluster; j++) {
//            sprintf(tmp_str, "%s %d", tmp_str, (int) label_map[j].size());
//        }
//        printf("%s\n", tmp_str);

        // get the distance to each cluster for each item and pass through the softmax
#pragma omp parallel for
        for (int j = 0; j < n_item; j++) {
            int *tmp_label_l = label_l + j * n_overlap;
            get_weight_l(base, tmp_label_l, label_map, dimension, n_cluster, j, weight_l[j]);
        }

//        printf("%.3f %.3f %.3f %.3f %.3f\n", weight_l[0][0], weight_l[0][1], weight_l[0][2], weight_l[0][3],
//               weight_l[0][4]);
        char save_path[256];
        sprintf(save_path, "%s/Classifier_%d/dataset_partition/weight_l.fvecs", basic_dir, i);
        save_weight_l(save_path, weight_l);
    }


    return 0;
}